import os, sys, traceback
import getpass
from threading import Thread
from subprocess import *
import numpy as np

if(sys.hexversion < 0x03000000):
	import Queue
else:
	import queue as Queue

from Searcher import Searcher

class GridSearcher(Searcher):
    """search best parameters using grid search"""
    telnet_workers = []
    ssh_workers = []
    nr_local_worker = 1
    def __init__(self, dataDisplayer, cmdline, *ranges):
        super(GridSearcher, self).__init__()
        self.dataDisplayer = dataDisplayer
        self.ranges = ranges
        self.cmdline = cmdline

    def search(self):
        """docstring for search"""
        self.jobs = self.calculate_jobs()
        #jobsfile = open('jobsfile', 'w')
        #jobsfile.write(str(self.jobs))
        self.job_queue = Queue.Queue(0)
        self.result_queue = Queue.Queue(0)

        for line in self.jobs:
            for (c,g) in line:
                self.job_queue.put((c,g))

        self.job_queue._put = self.job_queue.queue.appendleft


        # fire telnet workers

        if self.telnet_workers:
            nr_telnet_worker = len(self.telnet_workers)
            username = getpass.getuser()
            password = getpass.getpass()
            for host in self.telnet_workers:
                TelnetWorker(host,job_queue,self.result_queue,self.cmdline,
                         host,username,password).start()

        # fire ssh workers

        if self.ssh_workers:
            for host in self.ssh_workers:
                SSHWorker(host,self.job_queue,self.result_queue,self.cmdline,host).start()

        # fire local workers

        for i in range(self.nr_local_worker):
            LocalWorker('local',self.job_queue,self.result_queue,self.cmdline).start()

        # gather results

        done_jobs = {}
        db = []
        best_rate = -1
        best_c1,best_g1 = None,None
        self.row = None
        self.col = None
        for line in self.jobs:
            for (c,g) in line:
                while (c, g) not in done_jobs:
                    (worker,c1,g1,rate) = self.result_queue.get()
                    done_jobs[(c1,g1)] = rate
                    self.row = self.g_seq.index(g1)
                    self.col = self.c_seq.index(c1)
                    self.grid_rate[self.row,self.col] = rate
                    if (rate > best_rate) or (rate==best_rate and g1==best_g1 and c1<best_c1):
                        best_rate = rate
                        best_c1,best_g1=c1,g1
                        best_c = 2.0**c1
                        best_g = 2.0**g1
                    print("[%s] %s %s %s (best c=%s, g=%s, rate=%s)" % \
                (worker,c1,g1,rate, best_c, best_g, best_rate))
                    #self.dataDisplayer.writeConsole("[%s] %s %s %s (best c=%s, g=%s, rate=%s)" % \
                #(worker,c1,g1,rate, best_c, best_g, best_rate))
            #if len(line) > 0:
                ##x = self.grid_c[0:self.row+1,0:self.col+1]
                ##y = self.grid_g[0:self.row+1,0:self.col+1]
                ##z = self.grid_rate[0:self.row+1,0:self.col+1]
                ##self.dataDisplayer.drawContour(x,y,z)
                #self.dataDisplayer.drawContour(self.grid_c,self.grid_g,self.grid_rate)


        self.job_queue.put((WorkerStopToken,None))
        print("%s %s %s" % (best_c, best_g, best_rate))
        self.dataDisplayer.writeConsole("%s %s %s" % (best_c, best_g, best_rate))
        #self.dataDisplayer.drawContour(self.grid_c,self.grid_g,self.grid_rate)
        return (best_c, best_g)


    def range_f(self, begin,end,step):
        # like range, but works on non-integer too
        seq = []
        while True:
            if step > 0 and begin > end: break
            if step < 0 and begin < end: break
            seq.append(begin)
            begin = begin + step
        return seq

    def permute_sequence(self, seq):
        n = len(seq)
        if n <= 1: return seq

        mid = int(n/2)
        left = self.permute_sequence(seq[:mid])
        right = self.permute_sequence(seq[mid+1:])

        ret = [seq[mid]]
        while left or right:
            if left: ret.append(left.pop(0))
            if right: ret.append(right.pop(0))

        return ret
    def calculate_jobs(self):
        c_begin, c_end, c_step = self.ranges[0]
        g_begin, g_end, g_step = self.ranges[1]
        self.c_seq = self.permute_sequence(self.range_f(c_begin, c_end, c_step))
        self.g_seq = self.permute_sequence(self.range_f(g_begin, g_end, g_step))
        self.grid_c, self.grid_g = np.meshgrid(self.c_seq, self.g_seq)
        self.grid_rate = np.zeros(self.grid_c.shape, dtype=self.grid_c.dtype)
        print self.grid_c
        print self.grid_g
        print self.grid_rate
        nr_c = float(len(self.c_seq))
        nr_g = float(len(self.g_seq))
        i = 0
        j = 0
        jobs = []

        while i < nr_c or j < nr_g:
            if i/nr_c < j/nr_g:
                # increase C resolution
                line = []
                for k in range(0,j):
                    line.append((self.c_seq[i],self.g_seq[k]))
                i = i + 1
                jobs.append(line)
            else:
                # increase g resolution
                line = []
                for k in range(0,i):
                    line.append((self.c_seq[k],self.g_seq[j]))
                j = j + 1
                jobs.append(line)
        return jobs

class WorkerStopToken:  # used to notify the worker to stop
        pass

class Worker(Thread):
    def __init__(self,name,job_queue,result_queue, cmdline):
        Thread.__init__(self)
        self.name = name
        self.job_queue = job_queue
        self.result_queue = result_queue
        self.cmdline = cmdline
    def run(self):
        while True:
            (cexp,gexp) = self.job_queue.get()
            if cexp is WorkerStopToken:
                self.job_queue.put((cexp,gexp))
                # print 'worker %s stop.' % self.name
                break
            try:
                rate = self.run_one(2.0**cexp,2.0**gexp)
                if rate is None: raise RuntimeError("get no rate")
            except:
                # we failed, let others do that and we just quit

                traceback.print_exception(sys.exc_info()[0], sys.exc_info()[1], sys.exc_info()[2])

                self.job_queue.put((cexp,gexp))
                print('worker %s quit.' % self.name)
                break
            else:
                self.result_queue.put((self.name,cexp,gexp,rate))

class LocalWorker(Worker):
    def run_one(self,c,g):
        #cmdline = '%s -c %s -g %s -v %s %s %s' % (svmtrain_exe,c,g,fold,pass_through_string,dataset_pathname)
        cmdline = self.cmdline
        cmdline = cmdline.replace('cost', str(c))
        cmdline = cmdline.replace('gamma', str(g))
        result = Popen(cmdline,shell=True,stdout=PIPE).stdout
        for line in result.readlines():
            if str(line).find("Cross") != -1:
                return float(line.split()[-1][0:-1])

class SSHWorker(Worker):
    def __init__(self,name,job_queue,result_queue,cmdline,host):
        Worker.__init__(self,name,job_queue,result_queue,cmdline)
        self.host = host
        self.cwd = os.getcwd()
    def run_one(self,c,g):
        cmdline = 'ssh -x %s "cd %s;' % (self.host, self.cwd) + self.cmdline
        result = Popen(cmdline,shell=True,stdout=PIPE).stdout
        for line in result.readlines():
            if str(line).find("Cross") != -1:
                return float(line.split()[-1][0:-1])

class TelnetWorker(Worker):
    def __init__(self,name,job_queue,result_queue,cmdline,host,username,password):
        Worker.__init__(self,name,job_queue,result_queue,cmdline)
        self.host = host
        self.username = username
        self.password = password
    def run(self):
        import telnetlib
        self.tn = tn = telnetlib.Telnet(self.host)
        tn.read_until("login: ")
        tn.write(self.username + "\n")
        tn.read_until("Password: ")
        tn.write(self.password + "\n")

        # XXX: how to know whether login is successful?
        tn.read_until(self.username)
        #
        print('login ok', self.host)
        tn.write("cd "+os.getcwd()+"\n")
        Worker.run(self)
        tn.write("exit\n")
    def run_one(self,c,g):
        cmdline = self.cmdline
        result = self.tn.write(cmdline+'\n')
        (idx,matchm,output) = self.tn.expect(['Cross.*\n'])
        for line in output.split('\n'):
            if str(line).find("Cross") != -1:
                return float(line.split()[-1][0:-1])
